package control

import (
	"errors"
	"fmt"
	"github.com/xlvector/dlog"
	"strings"
)

type Message struct {
	Addr     string `json:"addr"`
	Ports    []int  `json:"ports"`
	Type     string `json:"type"`     //http https socks4 socks5
	Classify string `json:"classify"` //A B C H
}

func (p *Message) Check() error {
	/*
		switch p.Classify {
		case "A":
		case "B":
		case "C":
		case "H":
		default:
			return errors.New("classify wrong!")
		}
	*/

	switch p.Type {
	case "https":
	case "http":
	case "socks5":
	case "socks4":
	default:
		return errors.New("type wrong!")
	}

	return nil
}

func (p *Worker) workMsgA(msg *Message) {
	cnt := strings.Count(msg.Addr, ".")
	if cnt > 1 {
		dlog.Warn("task is not classify A! %v", msg)
		return
	}
	if len(msg.Ports) == 0 {
		dlog.Warn("task has not ports! %v", msg)
		return
	}

	baseAddr := msg.Addr
	if cnt == 0 {
		baseAddr += "."
	}

	for _, port := range msg.Ports {
		for i := 254; i > 0; i-- {
			for j := 255; j > 0; j-- {
				for k := 255; k > 0; k-- {
					p.taskch <- fmt.Sprintf("%s://%s%d.%d.%d:%d",
						msg.Type, baseAddr, k, j, i, port)
				}

			}
		}
	}

	p.stats.Add(SCAN_A_MSG, 1)
}

func (p *Worker) workMsgB(msg *Message) {
	cnt := strings.Count(msg.Addr, ".")
	if cnt > 2 || cnt < 1 {
		dlog.Warn("task is not classify B! %v", msg)
		return
	}
	if len(msg.Ports) == 0 {
		dlog.Warn("task has not ports! %v", msg)
		return
	}

	baseAddr := msg.Addr
	if cnt == 1 {
		baseAddr += "."
	}

	for _, port := range msg.Ports {
		for i := 254; i > 0; i-- {
			for j := 255; j > 0; j-- {
				p.taskch <- fmt.Sprintf("%s://%s%d.%d:%d",
					msg.Type, baseAddr, j, i, port)
			}
		}
	}
	p.stats.Add(SCAN_B_MSG, 1)
}

func (p *Worker) workMsgC(msg *Message) {
	cnt := strings.Count(msg.Addr, ".")
	if cnt > 3 || cnt < 2 {
		dlog.Warn("task is not classify C! %v", msg)
		return
	}
	if len(msg.Ports) == 0 {
		dlog.Warn("task has not ports! %v", msg)
		return
	}

	baseAddr := msg.Addr
	if cnt == 2 {
		baseAddr += "."
	} else if cnt == 3 {
		if !strings.HasSuffix(baseAddr, ".") {
			dlog.Warn("Msg C addr fail! %v", msg)
		}
	}

	for _, port := range msg.Ports {
		for i := 254; i > 0; i-- {
			p.taskch <- fmt.Sprintf("%s://%s%d:%d",
				msg.Type, baseAddr, i, port)
		}
	}
	p.stats.Add(SCAN_C_MSG, 1)
}

func (p *Worker) workMsgH(msg *Message) {
	cnt := strings.Count(msg.Addr, ".")
	if cnt != 3 || strings.HasSuffix(msg.Addr, ".") {
		dlog.Warn("task is not classify H! %v", msg)
		return
	}
	if len(msg.Ports) == 0 {
		dlog.Warn("task has not ports! %v", msg)
		return
	}

	baseAddr := msg.Addr
	if cnt == 2 {
		baseAddr += "."
	}

	for _, port := range msg.Ports {
		p.taskch <- fmt.Sprintf("%s://%s:%d",
			msg.Type, baseAddr, port)
	}
	p.stats.Add(SCAN_H_MSG, 1)
}

func (p *Worker) workMsg(idx int) {
	defer p.wg.Done()
	for {
		msg, ok := <-p.msgch
		if !ok {
			dlog.Warn("msgch exit! %d", idx)
			return
		}

		//check msg
		if err := msg.Check(); err != nil {
			dlog.Warn("msg check fail! %s", err)
			continue
		}

		switch msg.Classify {
		case "A":
			p.workMsgA(msg)
		case "B":
			p.workMsgB(msg)
		case "C":
			p.workMsgC(msg)
		case "H":
			p.workMsgH(msg)
		default:
			dlog.Warn("classify wrong! %v", msg)
		}
	}
}
